
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, \
    FewShotChatMessagePromptTemplate
from langchain_openai import ChatOpenAI
from env_utils import LOCAL_API_KEY, LOCAL_BASE_URL
from langchain_core.messages import HumanMessage, SystemMessage

# llm = ChatOpenAI(
#     model="deepseek-chat",
#     temperature=0.8,
#     api_key=DEEPSEEK_API_KEY,
#     base_url=DEEPSEEK_BASE_URL
# )

llm = ChatOpenAI(
    model="qwen3-8b",
    temperature=0.8,
    api_key=LOCAL_API_KEY,
    base_url=LOCAL_BASE_URL,
    extra_body={'chat_template_kwargs': {'enable_thinking': True}},
)

examples = [
    {"input": "1$$2", "output": "3"},
    {"input": "4$$5", "output": "9"}
]


# messages = [SystemMessage('你是一个资深动漫专家，你了解所有的动漫信息'),
#             HumanMessage('你知道高达SEED吗？')]
# resp = llm.invoke(messages)

base_prompt = ChatPromptTemplate.from_messages([
    ('human', '{input}'),
    ('ai', '{output}')
])
#包含实例的提示词模板
few_shot_prompt = FewShotChatMessagePromptTemplate(
    examples=examples,
    example_prompt=base_prompt,
)

final_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个智能助手"),
    few_shot_prompt,
    MessagesPlaceholder("msgs")
])

chain = final_prompt | llm
resp = chain.invoke({"msgs": [HumanMessage("3$$3")]})
print(resp)


